"""
增强的IP频率限制中间件
支持Redis持久化存储、多层级限制、IP白名单/黑名单等功能
"""

import asyncio
import json
import time
from typing import Dict, List, Optional, Set, Union
from fastapi import Request, status
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
import redis.asyncio as redis
from bot_api_v1.app.core.config import settings
from bot_api_v1.app.core.logger import logger


class EnhancedRateLimitMiddleware(BaseHTTPMiddleware):
    """
    增强的IP频率限制中间件
    
    功能特性：
    - Redis持久化存储
    - 多层级频率限制（分钟/小时/日）
    - IP白名单/黑名单
    - 地理位置感知（可选）
    - 实时监控和报警
    - API Key级别的差异化限制
    """
    
    def __init__(self, app):
        super().__init__(app)
        self.redis_client: Optional[redis.Redis] = None
        self.fallback_storage: Dict = {}  # 在Redis不可用时的备用存储
        
        # 配置参数
        self.config = {
            # 基础频率限制（每分钟）
            'base_limit_per_minute': getattr(settings, 'RATE_LIMIT_PER_MINUTE', 60),
            'base_limit_per_hour': getattr(settings, 'RATE_LIMIT_PER_HOUR', 1000), 
            'base_limit_per_day': getattr(settings, 'RATE_LIMIT_PER_DAY', 10000),
            
            # VIP用户限制（更高限制）
            'vip_limit_per_minute': getattr(settings, 'VIP_RATE_LIMIT_PER_MINUTE', 120),
            'vip_limit_per_hour': getattr(settings, 'VIP_RATE_LIMIT_PER_HOUR', 3000),
            'vip_limit_per_day': getattr(settings, 'VIP_RATE_LIMIT_PER_DAY', 50000),
            
            # 严格模式限制（对可疑IP）
            'strict_limit_per_minute': getattr(settings, 'STRICT_RATE_LIMIT_PER_MINUTE', 10),
            'strict_limit_per_hour': getattr(settings, 'STRICT_RATE_LIMIT_PER_HOUR', 100),
            'strict_limit_per_day': getattr(settings, 'STRICT_RATE_LIMIT_PER_DAY', 500),
            
            # 窗口大小（秒）
            'minute_window': 60,
            'hour_window': 3600,
            'day_window': 86400,
            
            # Redis配置
            'redis_key_prefix': 'rate_limit',
            'redis_expire_time': 86400,  # 24小时
            
            # IP列表配置
            'whitelist_enabled': getattr(settings, 'IP_WHITELIST_ENABLED', True),
            'blacklist_enabled': getattr(settings, 'IP_BLACKLIST_ENABLED', True),
        }
        
        # IP白名单（永不限制的IP）
        self.ip_whitelist: Set[str] = self._load_ip_whitelist()
        
        # IP黑名单（完全阻止的IP）
        self.ip_blacklist: Set[str] = self._load_ip_blacklist()
        
        # 可疑IP列表（使用严格限制）
        self.suspicious_ips: Set[str] = set()
        
        # 初始化Redis连接
        asyncio.create_task(self._init_redis())
    
    async def _init_redis(self) -> None:
        """初始化Redis连接"""
        try:
            redis_url = getattr(settings, 'REDIS_URL', 'redis://localhost:6379')
            self.redis_client = redis.from_url(
                redis_url,
                encoding='utf-8',
                decode_responses=True,
                socket_connect_timeout=5,
                socket_keepalive=True,
                health_check_interval=30
            )
            
            # 测试连接
            await self.redis_client.ping()
            logger.info("Enhanced Rate Limit Middleware: Redis connection established")
            
        except Exception as e:
            logger.warning(f"Enhanced Rate Limit Middleware: Redis connection failed: {e}")
            logger.warning("Falling back to in-memory storage")
            self.redis_client = None
    
    def _load_ip_whitelist(self) -> Set[str]:
        """加载IP白名单"""
        whitelist = set()
        
        # 从配置文件加载
        config_whitelist = getattr(settings, 'IP_WHITELIST', [])
        if isinstance(config_whitelist, (list, tuple)):
            whitelist.update(config_whitelist)
        
        # 默认添加本地IP
        whitelist.update(['127.0.0.1', '::1', 'localhost'])
        
        logger.info(f"Loaded IP whitelist: {len(whitelist)} entries")
        return whitelist
    
    def _load_ip_blacklist(self) -> Set[str]:
        """加载IP黑名单"""
        blacklist = set()
        
        # 从配置文件加载
        config_blacklist = getattr(settings, 'IP_BLACKLIST', [])
        if isinstance(config_blacklist, (list, tuple)):
            blacklist.update(config_blacklist)
        
        logger.info(f"Loaded IP blacklist: {len(blacklist)} entries")
        return blacklist
    
    async def dispatch(self, request: Request, call_next):
        """主要的请求处理逻辑"""
        # 获取客户端IP
        client_ip = self._get_client_ip(request)
        
        # 检查黑名单
        if self.config['blacklist_enabled'] and client_ip in self.ip_blacklist:
            # 黑名单阻止是严重安全事件，使用 error 级别自动入库
            logger.error(
                f"黑名单IP访问被阻止: {client_ip} 尝试访问 {request.method} {request.url.path}",
                extra={
                    "event_type": "blacklist_ip_blocked",
                    "client_ip": client_ip,
                    "action": "blocked",
                    "reason": "blacklisted_ip", 
                    "path": str(request.url.path),
                    "method": request.method,
                    "user_agent": request.headers.get('User-Agent', 'Unknown'),
                    "x_forwarded_for": request.headers.get("X-Forwarded-For"),
                    "x_real_ip": request.headers.get("X-Real-IP"),
                    "referer": request.headers.get("Referer")
                }
            )
            
            return self._create_blocked_response("IP blocked")
        
        # 检查白名单
        if self.config['whitelist_enabled'] and client_ip in self.ip_whitelist:
            # 白名单IP直接通过，不进行限制
            return await call_next(request)
        
        # 获取用户类型
        user_type = await self._determine_user_type(request)
        
        # 检查频率限制
        is_limited, limit_info = await self._check_rate_limits(client_ip, user_type)
        
        if is_limited:
            # 记录限制事件
            await self._log_rate_limit_event(client_ip, request, limit_info)
            
            # 检查是否需要将IP标记为可疑
            await self._check_suspicious_behavior(client_ip)
            
            return self._create_rate_limit_response(limit_info)
        
        # 记录请求
        await self._record_request(client_ip, user_type)
        
        # 处理请求
        try:
            response = await call_next(request)
            
            # 可选：基于响应状态进行后处理
            await self._post_process_response(client_ip, response)
            
            return response
            
        except Exception as e:
            logger.error(f"Error in rate limit middleware: {e}")
            return await call_next(request)
    
    def _get_client_ip(self, request: Request) -> str:
        """获取真实客户端IP"""
        # 检查常见的代理头
        forwarded_ips = [
            request.headers.get("X-Forwarded-For"),
            request.headers.get("X-Real-IP"),
            request.headers.get("CF-Connecting-IP"),  # Cloudflare
            request.headers.get("X-Client-IP"),
        ]
        
        for ip in forwarded_ips:
            if ip:
                # X-Forwarded-For 可能包含多个IP，取第一个
                return ip.split(',')[0].strip()
        
        # 回退到直接连接IP
        return request.client.host if request.client else "unknown"
    
    async def _determine_user_type(self, request: Request) -> str:
        """确定用户类型以应用不同的限制策略"""
        # 检查API Key以确定用户级别
        auth_header = request.headers.get("Authorization") or request.headers.get("X-Auth-Key")
        
        if auth_header:
            # 这里可以集成数据库查询来确定用户的VIP状态
            # 暂时基于API Key存在性判断
            return "authenticated"
        
        # 检查是否为可疑IP
        client_ip = self._get_client_ip(request)
        if client_ip in self.suspicious_ips:
            return "suspicious"
        
        return "anonymous"
    
    async def _check_rate_limits(self, client_ip: str, user_type: str) -> tuple[bool, dict]:
        """检查多层级频率限制"""
        current_time = time.time()
        windows = ['minute', 'hour', 'day']
        
        for window in windows:
            window_size = self.config[f'{window}_window']
            limit_key = f"{user_type}_limit_per_{window}"
            
            if user_type == "suspicious":
                limit_key = f"strict_limit_per_{window}"
            elif user_type == "authenticated":
                # 可以根据具体API Key确定是否为VIP
                limit_key = f"base_limit_per_{window}"
            else:  # anonymous
                limit_key = f"base_limit_per_{window}"
            
            limit = self.config.get(limit_key, self.config[f'base_limit_per_{window}'])
            
            # 检查当前窗口的请求数量
            request_count = await self._get_request_count(client_ip, window, window_size, current_time)
            
            if request_count >= limit:
                return True, {
                    'window': window,
                    'limit': limit,
                    'current': request_count,
                    'reset_time': current_time + window_size,
                    'user_type': user_type
                }
        
        return False, {}
    
    async def _get_request_count(self, client_ip: str, window: str, window_size: int, current_time: float) -> int:
        """获取指定时间窗口内的请求数量"""
        redis_key = f"{self.config['redis_key_prefix']}:{client_ip}:{window}"
        
        try:
            if self.redis_client:
                # 使用Redis ZCOUNT来计算时间窗口内的请求数
                min_time = current_time - window_size
                count = await self.redis_client.zcount(redis_key, min_time, current_time)
                return count
            else:
                # 回退到内存存储
                requests = self.fallback_storage.get(redis_key, [])
                # 过滤出时间窗口内的请求
                min_time = current_time - window_size
                valid_requests = [t for t in requests if t >= min_time]
                self.fallback_storage[redis_key] = valid_requests
                return len(valid_requests)
                
        except Exception as e:
            logger.error(f"Error getting request count: {e}")
            return 0
    
    async def _record_request(self, client_ip: str, user_type: str) -> None:
        """记录请求到存储"""
        current_time = time.time()
        windows = ['minute', 'hour', 'day']
        
        try:
            if self.redis_client:
                # 使用Redis有序集合记录请求时间戳
                for window in windows:
                    redis_key = f"{self.config['redis_key_prefix']}:{client_ip}:{window}"
                    
                    # 添加当前请求时间戳
                    await self.redis_client.zadd(redis_key, {str(current_time): current_time})
                    
                    # 设置过期时间
                    window_size = self.config[f'{window}_window']
                    await self.redis_client.expire(redis_key, window_size * 2)
                    
                    # 清理过期的条目
                    min_time = current_time - window_size
                    await self.redis_client.zremrangebyscore(redis_key, 0, min_time)
            else:
                # 回退到内存存储
                for window in windows:
                    redis_key = f"{self.config['redis_key_prefix']}:{client_ip}:{window}"
                    
                    if redis_key not in self.fallback_storage:
                        self.fallback_storage[redis_key] = []
                    
                    self.fallback_storage[redis_key].append(current_time)
                    
                    # 清理过期条目
                    window_size = self.config[f'{window}_window']
                    min_time = current_time - window_size
                    self.fallback_storage[redis_key] = [
                        t for t in self.fallback_storage[redis_key] if t >= min_time
                    ]
                    
        except Exception as e:
            logger.error(f"Error recording request: {e}")
    
    async def _check_suspicious_behavior(self, client_ip: str) -> None:
        """检查并标记可疑行为"""
        try:
            # 检查短时间内的高频请求
            current_time = time.time()
            short_window = 10  # 10秒窗口
            request_count = await self._get_request_count(client_ip, 'minute', short_window, current_time)
            
            # 如果10秒内超过30个请求，标记为可疑
            if request_count > 30:
                self.suspicious_ips.add(client_ip)
                
                # 可疑行为检测记录到数据库，使用 info_to_db
                logger.info_to_db(
                    f"检测到可疑IP行为: {client_ip} 在{short_window}秒内发送了{request_count}个请求，已标记为可疑",
                    extra={
                        "event_type": "suspicious_ip_detected",
                        "client_ip": client_ip,
                        "requests_in_10s": request_count,
                        "detection_threshold": 30,
                        "detection_window": short_window,
                        "detection_reason": "high_frequency_requests",
                        "action": "marked_suspicious"
                    }
                )
                
                # 可以在这里添加更多的检测逻辑，如：
                # - 检查User-Agent模式
                # - 检查请求路径模式  
                # - 地理位置异常检测等
                
        except Exception as e:
            logger.error(f"Error checking suspicious behavior: {e}")
    
    async def _log_rate_limit_event(self, client_ip: str, request: Request, limit_info: dict) -> None:
        """记录频率限制事件"""
        event_data = {
            'client_ip': client_ip,
            'path': str(request.url.path),
            'method': request.method,
            'user_agent': request.headers.get('User-Agent', 'Unknown'),
            'limit_info': limit_info,
            'timestamp': time.time()
        }
        
        # 频率限制触发记录到数据库，使用 info_to_db
        logger.info_to_db(
            f"频率限制触发: IP {client_ip} 在 {limit_info['window']} 窗口内请求 {limit_info['current']}/{limit_info['limit']} 次",
            extra={
                "event_type": "rate_limit_exceeded",
                "client_ip": client_ip,
                "window_type": limit_info['window'],
                "limit": limit_info['limit'],
                "current_requests": limit_info['current'],
                "user_type": limit_info.get('user_type', 'unknown'),
                "reset_time": limit_info.get('reset_time'),
                "path": str(request.url.path),
                "method": request.method,
                "user_agent": request.headers.get('User-Agent', 'Unknown'),
                "x_forwarded_for": request.headers.get("X-Forwarded-For"),
                "x_real_ip": request.headers.get("X-Real-IP"),
                "referer": request.headers.get("Referer")
            }
        )
    
    async def _post_process_response(self, client_ip: str, response) -> None:
        """响应后处理"""
        # 可以根据响应状态码进行额外的处理
        # 例如：4xx/5xx响应可能表示恶意请求
        if hasattr(response, 'status_code'):
            if response.status_code >= 400:
                # 记录错误响应，可用于后续分析
                pass
    
    def _create_rate_limit_response(self, limit_info: dict) -> JSONResponse:
        """创建频率限制响应"""
        retry_after = int(limit_info.get('reset_time', time.time() + 60) - time.time())
        
        return JSONResponse(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            content={
                "code": 429,
                "message": f"Rate limit exceeded for {limit_info['window']} window",
                "data": {
                    "limit": limit_info['limit'],
                    "current": limit_info['current'],
                    "window": limit_info['window'],
                    "retry_after": max(retry_after, 1),
                    "user_type": limit_info.get('user_type', 'unknown')
                }
            },
            headers={
                "X-RateLimit-Limit": str(limit_info['limit']),
                "X-RateLimit-Remaining": str(max(0, limit_info['limit'] - limit_info['current'])),
                "X-RateLimit-Reset": str(int(limit_info['reset_time'])),
                "Retry-After": str(max(retry_after, 1))
            }
        )
    
    def _create_blocked_response(self, reason: str) -> JSONResponse:
        """创建阻止访问响应"""
        return JSONResponse(
            status_code=status.HTTP_403_FORBIDDEN,
            content={
                "code": 403,
                "message": "Access denied",
                "data": {
                    "reason": reason,
                    "timestamp": int(time.time())
                }
            }
        )
    
    async def get_rate_limit_stats(self, client_ip: str) -> dict:
        """获取指定IP的频率限制统计信息（用于调试和监控）"""
        current_time = time.time()
        stats = {}
        
        for window in ['minute', 'hour', 'day']:
            window_size = self.config[f'{window}_window']
            count = await self._get_request_count(client_ip, window, window_size, current_time)
            
            stats[window] = {
                'count': count,
                'window_size': window_size,
                'remaining': max(0, self.config[f'base_limit_per_{window}'] - count)
            }
        
        return {
            'ip': client_ip,
            'stats': stats,
            'is_whitelisted': client_ip in self.ip_whitelist,
            'is_blacklisted': client_ip in self.ip_blacklist,
            'is_suspicious': client_ip in self.suspicious_ips,
            'timestamp': current_time
        }